// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu, jhkim

#include "coin2/strategy/linear_model/agg2.h"

#include <algorithm>

#include "coin2/strategy/linear_model/util.h"

namespace coin2::strategy::linear_model {

/* ------------------------------ Order Protos ------------------------------ */
using coin::proto::OrderDirection;
using coin::proto::OrderDuration;
using coin::proto::OrderType;


void LmAggLogic2::onFeed(
    const Symbol* symbol,
    const FeedUpdate& upd,
    const ::fastfeature::FeedMessage& feed_msg,
    LinearModel* linear_model,
    int64_t* lm_feed_id,
    const int64_t* global_feed_id,
    LinearModelDelegateProto* delegate,
    bool simulation) {
  int64_t timestamp_override = upd.timestamp();
  std::optional<double> midp_return;
  lholder_.UpdateFeed(symbol, upd, feed_msg);
  if (!lholder_.HasBookAndPos()) {
    SPDLOG_INFO_EVERY_NS(
        10'000'000'000LL,
        timestamp_override,
        "{}: BookAndPos is not available yet.",
        symbol->symbol_string());
    return;
  }
  filldepth_window_.UpdateTimeOnly(timestamp_override);

  if (upd.product() == lholder_.target_product && upd.is_trade()) {
    const Trade& trade = upd.trade();
    filldepth_window_.UpdateTradep(timestamp_override, trade);
  } else if (upd.product() == lholder_.target_product && upd.is_book()) {
    filldepth_window_.UpdateBookp(timestamp_override, lholder_.last_book);
  }

  TryTrigger(
      symbol,
      upd,
      feed_msg,
      linear_model,
      lm_feed_id,
      global_feed_id,
      delegate,
      timestamp_override,
      &midp_return);

  bool use_model = lm_agg_config_.filldepth_window_use_pred_model();
  int use_simple = lm_agg_config_.filldepth_window_use_simple_ver();
  std::optional<double> buy_filldepth_pred;
  std::optional<double> sell_filldepth_pred;
  if ((use_model || use_simple) && !delegate->response_.buy_filldepth_model().empty()) {
    buy_filldepth_pred = delegate->response_.buy_filldepth_model().Get(0).signal();
  }
  if ((use_model || use_simple) && !delegate->response_.sell_filldepth_model().empty()) {
    sell_filldepth_pred = delegate->response_.sell_filldepth_model().Get(0).signal();
  }

  if (midp_return && lm_agg_config_.use_cancel_model()) {
    filldepth_window_.UpdateSignal(timestamp_override, *midp_return);
    // *midp_return *= lm_agg_config_.signal_multiplier();
    auto ask0p = lholder_.last_book->Ask0()->price;
    auto bid0p = lholder_.last_book->Bid0()->price;
    double midp = 0.5 * (ask0p + bid0p);

    double signal_lb;
    double signal_ub;
    if (!pred_interval_.GetNSigmaBound(
            *midp_return,
            lm_agg_config_.nsigma_edge(),
            &signal_lb,
            &signal_ub,
            lm_agg_config_.nsigma_edge_cap())) {
      return;
    }

    // bool use_sh = lm_agg_config_.filldepth_window_use_sh_signal();
    // std::optional<double> sh_signal;
    // if (use_sh && !delegate->response_.buy_sh_model().empty()) {
    //   sh_signal = delegate->response_.buy_sh_model().Get(0).signal();
    // }

    double ask_skew_edge_bps = 0;
    double bid_skew_edge_bps = 0;
    if (lm_agg_config_.filldepth_window_use_edge_clearing()) {
      double current_stack = lholder_.oe->GetCurrentStack(*lholder_.product_holder);

      double skew_edge_bps = 0;
      clear_edge_manager_.UpdatePositionInStack(timestamp_override, current_stack);
      skew_edge_bps = clear_edge_manager_.clear_lean_bps
          - current_stack * lm_agg_config_.lean_per_stack_bps();

      ask_skew_edge_bps = std::max(skew_edge_bps, -5.);
      signal_ub += ask_skew_edge_bps;

      bid_skew_edge_bps = std::min(skew_edge_bps, 5.);
      signal_lb += bid_skew_edge_bps;
    }

    double buy_edge = -1001;
    double sell_edge = 1001;
    double buy_qty_multiplier = 1;
    double sell_qty_multiplier = 1;
    bool edge_available = false;

    edge_available = filldepth_window_.GetEdge(
        timestamp_override,
        *midp_return,
        &signal_lb,
        &signal_ub,
        ask_skew_edge_bps,
        bid_skew_edge_bps,
        &lm_agg_config_,
        &buy_filldepth_pred,
        &sell_filldepth_pred,
        &buy_edge,
        &sell_edge,
        &buy_qty_multiplier,
        &sell_qty_multiplier);

    if (!edge_available) {
      SPDLOG_INFO_EVERY_NS(
          10'000'000'000LL,
          timestamp_override,
          "{}: edge is not available yet, rest..",
          symbol->symbol_string());
      rest_mode_ = true;
      TryClear(symbol, upd, feed_msg);
      return;
    } else {
      rest_mode_ = false;
    }

    double cancel_edge_bps = lm_agg_config_.cancel_edge_bps();
    double cancel_outside_bps = lm_agg_config_.cancel_outside_bps();

    const auto* ask_inside_order =
        lholder_.oe->GetInsideOrder(*lholder_.product_holder, OrderDirection::DIRECTION_SELL);
    double ask_post_price = ask0p * (1 + 1e-4 * sell_edge);
    double ask_inside_bound = ask_post_price * (1 - 1e-4 * cancel_edge_bps);
    double ask_outside_bound = ask_post_price * (1 + 1e-4 * cancel_outside_bps);

    if (use_simple) {
      ask_inside_bound = midp * (1 + 1e-4 * (signal_ub + cancel_edge_bps_));
      ask_outside_bound = ask0p * (1 + 1e-4 * lm_agg_config_.cancel_outside_bps());
      ask_post_price = midp * (1 + (*midp_return + sell_edge) * 1e-4);
      if (use_simple == 2) {
        ask_post_price = ask_inside_bound * (1 + sell_edge * 1e-4);
      }
    }

    if (midp > prev_ask_fill_price_ &&
        prev_ask_pfill_ts_ + def_edge_ns_ > timestamp_override) {
      if (new_ask_fill_) {
        sell_edge_add_ += lm_agg_config_.defensive_edge_add_bps();
        new_ask_fill_ = false;
      }
    } else {
      if (prev_sell_def_edge_reduced_ts_ + def_edge_ns_ < timestamp_override) {
        sell_edge_add_ -= lm_agg_config_.defensive_edge_sub_bps();
        sell_edge_add_ = std::max(sell_edge_add_, 0.);
        prev_sell_def_edge_reduced_ts_ = timestamp_override;
      }
    }

    ask_post_price *= (1 + 1e-4 * (sell_edge_add_));
    ask_inside_bound *= (1 + 1e-4 * (sell_edge_add_));
    ask_outside_bound *= (1 + 1e-4 * (sell_edge_add_));

    bool ask_post = edge_available && (sell_edge < 999);
    if (ask_inside_order) {
      double ask_working_price = ask_inside_order->order_price();
      // LOG(INFO) << "ask_working_price= " << ask_working_price;
      if (ask_working_price < ask_inside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_SELL, timestamp_override);
        ask_post = ask_post && lm_agg_config_.layering();
      } else if (ask_working_price > ask_outside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_SELL, timestamp_override);
        ask_post = ask_post && lm_agg_config_.layering();
      } else {
        ask_post = false;
      }
    }

    SPDLOG_INFO_EVERY_NS(
        10'000'000'000LL,
        timestamp_override,
        "{}: {}, {}, {}, {}, {}, {}",
        symbol->symbol_string(),
        edge_available,
        sell_edge,
        ask_post,
        lm_agg_config_.use_passive_post(),
        prev_ask_post_time_ + post_cooltime_ < timestamp_override,
        prev_ask_fill_ts_ + fill_cooltime_ < timestamp_override);

    if (ask_post &&
        lm_agg_config_.use_passive_post() &&
        prev_ask_post_time_ + post_cooltime_ < timestamp_override &&
        prev_ask_fill_ts_ + fill_cooltime_ < timestamp_override) {
      prev_ask_post_time_ = timestamp_override;
      if (lm_agg_config_.filldepth_window_price_bound()) {
        ask_post_price = std::max(
            ask_post_price,
            lholder_.product_holder->product_info().tick_price().GetNextPrice(bid0p));
        ask_post_price = std::max(
            ask_post_price,
            bid0p * (1 + lm_agg_config_.min_edge_bps() * 1e-4));
      }
      if (lm_agg_config_.filldepth_window_use_bbo()) {
        ask_post_price = std::max(
            ask_post_price,
            ask0p * (1 + lm_agg_config_.min_edge_bps() * 1e-4));
      }

      double rounded_price =
          lholder_.product_holder->product_info().tick_price().RoundPassively(
              ask_post_price, OrderDirection::DIRECTION_SELL);
      double raw_qty = lholder_.oe->GetConfig(*lholder_.product_holder).qty().max_lot_size();
      raw_qty *= lm_agg_config_.maker_qty_ratio();
      raw_qty *= sell_qty_multiplier;

      lholder_.oe->SubmitOrder(
          *lholder_.product_holder,
          rounded_price,  // OE does rounding, so this can be replaced by
                          // ExecutorConfig-RoundingPolicy
          raw_qty,
          OrderDirection::DIRECTION_SELL,
          lm_agg_config_.order_duration(),
          OrderType::LIMIT_ORDER,
          ask0p,
          leverage_rate_,
          true,  // post_only
          "passive_post",
          timestamp_override,
          false,
          lm_agg_config_.use_block_price_band());
      // LOG(INFO) << "sell order submitted";
    }


    const auto* bid_inside_order =
        lholder_.oe->GetInsideOrder(*lholder_.product_holder, OrderDirection::DIRECTION_BUY);
    double bid_post_price = bid0p * (1 + buy_edge * 1e-4);
    double bid_inside_bound = bid_post_price * (1 + 1e-4 * cancel_edge_bps);
    double bid_outside_bound = bid_post_price * (1 - 1e-4 * cancel_outside_bps);

    if (use_simple) {
      bid_inside_bound = midp * (1 + 1e-4 * (signal_lb - cancel_edge_bps_));
      bid_outside_bound = bid0p * (1 - 1e-4 * lm_agg_config_.cancel_outside_bps());
      bid_post_price = midp * (1 + (*midp_return + buy_edge) * 1e-4);
      if (use_simple == 2) {
        bid_post_price = bid_inside_bound * (1 + buy_edge * 1e-4);
      }
    }

    if (midp < prev_bid_fill_price_ &&
        prev_bid_pfill_ts_ + def_edge_ns_ > timestamp_override) {
      if (new_bid_fill_) {
        buy_edge_add_ += lm_agg_config_.defensive_edge_add_bps();
        new_bid_fill_ = false;
      }
    } else {
      if (prev_buy_def_edge_reduced_ts_ + def_edge_ns_ < timestamp_override) {
        buy_edge_add_ -= lm_agg_config_.defensive_edge_sub_bps();
        buy_edge_add_ = std::max(buy_edge_add_, 0.);
        prev_buy_def_edge_reduced_ts_ = timestamp_override;
      }
    }

    bid_post_price *= (1 - 1e-4 * (buy_edge_add_));
    bid_inside_bound *= (1 - 1e-4 * (buy_edge_add_));
    bid_outside_bound *= (1 - 1e-4 * (buy_edge_add_));

    bool bid_post = edge_available && (buy_edge > -999);
    if (bid_inside_order) {
      double bid_working_price = bid_inside_order->order_price();
      // LOG(INFO) << "bid_working_price= " << bid_working_price;
      if (bid_working_price > bid_inside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_BUY, timestamp_override);
        bid_post = bid_post && lm_agg_config_.layering();
        // bid_post = lm_agg_config_.layering();
      } else if (bid_working_price < bid_outside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_BUY, timestamp_override);
        bid_post = bid_post && lm_agg_config_.layering();
        // bid_post = lm_agg_config_.layering();
      } else {
        bid_post = false;
      }
    }
    if (bid_post &&
        lm_agg_config_.use_passive_post() &&
        prev_bid_post_time_ + post_cooltime_ < timestamp_override &&
        prev_bid_fill_ts_ + fill_cooltime_ < timestamp_override) {
      prev_bid_post_time_ = timestamp_override;
      if (lm_agg_config_.filldepth_window_price_bound()) {
        bid_post_price = std::min(
            bid_post_price,
            lholder_.product_holder->product_info().tick_price().GetPrevPrice(ask0p));
        bid_post_price = std::min(
            bid_post_price,
            ask0p * (1 - lm_agg_config_.min_edge_bps() * 1e-4));
      }
      if (lm_agg_config_.filldepth_window_use_bbo()) {
        bid_post_price = std::min(
            bid_post_price,
            bid0p * (1 - lm_agg_config_.min_edge_bps() * 1e-4));
      }

      double rounded_price =
          lholder_.product_holder->product_info().tick_price().RoundPassively(
              bid_post_price, OrderDirection::DIRECTION_BUY);
      double raw_qty = lholder_.oe->GetConfig(*lholder_.product_holder).qty().max_lot_size();
      raw_qty *= lm_agg_config_.maker_qty_ratio();
      raw_qty *= buy_qty_multiplier;

      lholder_.oe->SubmitOrder(
          *lholder_.product_holder,
          rounded_price,  // OE does rounding, so this can be replaced by
                          // ExecutorConfig-RoundingPolicy
          raw_qty,
          OrderDirection::DIRECTION_BUY,
          lm_agg_config_.order_duration(),
          OrderType::LIMIT_ORDER,
          bid0p,
          leverage_rate_,
          true,  // post_only
          "passive_post",
          timestamp_override,
          false,
          lm_agg_config_.use_block_price_band());
      // LOG(INFO) << "buy order submitted";
    }
  }
  lholder_.oe->onPeriodicOrderCheck(*lholder_.product_holder, timestamp_override);
  if (!simulation) {
    lholder_.LogSignal(upd.timestamp(), delegate);
  }
}


void LmAggLogic2::TryTrigger(
    const Symbol* symbol,
    const FeedUpdate& upd,
    const ::fastfeature::FeedMessage& feed_msg,
    LinearModel* linear_model,
    int64_t* lm_feed_id,
    const int64_t* global_feed_id,
    LinearModelDelegateProto* delegate,
    int64_t timestamp_override,
    std::optional<double>* midp_return) {
  CHECK(lholder_.target_symbol);
  delegate->ClearProto();
  std::vector<double> feature_values;
  std::vector<std::string> feature_keys;
  lholder_.TriggerModelWithLatencyRecording(linear_model, lholder_.target_symbol);

  /* -------------------------------------------------------------------------- */
  /*                               lm_agg_logic.py                              */
  /* -------------------------------------------------------------------------- */

  double buy_edge_mismatch_bps = 0;
  double sell_edge_mismatch_bps = 0;
  if (pos_mismatch_calculator_) {
    pos_mismatch_calculator_->UpdateTimeOnly(timestamp_override);
    double edge_mismatch_stack = pos_mismatch_calculator_->GetMismatchInStackCount();
    double mismatch_edge_per_stack = lm_agg_config_.edge_desired_pos_mismatch_per_stack_bps();
    if (edge_mismatch_stack < 0) {
      // current position is more negative than desired
      sell_edge_mismatch_bps = std::abs(edge_mismatch_stack) * mismatch_edge_per_stack;
    } else {
      // current position is more positive than desired
      buy_edge_mismatch_bps = std::abs(edge_mismatch_stack) * mismatch_edge_per_stack;
    }
  }

  auto ask0p = lholder_.last_book->Ask0()->price;
  auto bid0p = lholder_.last_book->Bid0()->price;

  /* ----------------------------- pick up signal ----------------------------- */
  double signal = 0;
  double sign = 0;
  double bbo_price = 0;
  double bbo_price_pass = 0;
  OrderDirection direction = OrderDirection::DIRECTION_UNINITIALIZED;
  std::string model_id = "";

  double buy_pie_in_the_sky = 0;
  double sell_pie_in_the_sky = 0;

  if (!delegate->response_.buy_edge_model().empty()) {
    buy_pie_in_the_sky = delegate->response_.buy_edge_model().Get(0).signal();
  }
  if (!delegate->response_.sell_edge_model().empty()) {
    sell_pie_in_the_sky = delegate->response_.sell_edge_model().Get(0).signal();
  }

  double pie_in_the_sky = 0;
  double leaned_threshold_bps = threshold_bps_;

  double calc_stack = lholder_.oe->GetCurrentStack(*lholder_.product_holder);

  CHECK(midp_model_);
  if (!delegate->response_.buy_model().empty()) {
    double midp_ret = delegate->response_.buy_model().Get(0).signal();
    *midp_return = midp_ret;
    double spread_bps = (ask0p - bid0p) / ask0p * 1e4;
    model_id = delegate->response_.buy_model().Get(0).tag();
    if (midp_ret > 0) {
      sign = 1;
      direction = OrderDirection::DIRECTION_BUY;
      bbo_price = ask0p;
      bbo_price_pass = bid0p;
      pie_in_the_sky = buy_pie_in_the_sky;
      signal = midp_ret - penalize_spread_ratio_ * spread_bps;
      leaned_threshold_bps = threshold_bps_ +
                             (calc_stack > 0) * calc_stack * lm_agg_config_.lean_per_stack_bps();
    } else {
      sign = -1;
      direction = OrderDirection::DIRECTION_SELL;
      bbo_price = bid0p;
      bbo_price_pass = ask0p;
      pie_in_the_sky = sell_pie_in_the_sky;
      signal = -midp_ret - penalize_spread_ratio_ * spread_bps;
      leaned_threshold_bps =
          threshold_bps_ +
          (calc_stack < 0) * (-calc_stack) * lm_agg_config_.lean_per_stack_bps();
    }
  }

  // signal window
  pred_interval_.UpdateSignal(timestamp_override, signal);
  double signal_lb;
  double signal_ub;
  if (!pred_interval_.GetNSigmaBound(
        signal,
        lm_agg_config_.nsigma_edge(),
        &signal_lb,
        &signal_ub,
        lm_agg_config_.nsigma_edge_cap())) {
    return;
  }

  // edesired
  double edge_mismatch_bps = 0;

  if (sign > 0) {
    edge_mismatch_bps = buy_edge_mismatch_bps;
  } else if (sign < 0) {
    edge_mismatch_bps = sell_edge_mismatch_bps;
  } else {
    return;
  }

  leaned_threshold_bps = std::min<double>(leaned_threshold_bps, threshold_bps_ * 2);

  leaned_threshold_bps += edge_mismatch_bps;

  /* ----------------------------- threshold check ---------------------------- */
  double midp = 0.5 * (ask0p + bid0p);
  // auto signal_to_bp_divisor = is_model_output_bps_ ? 1 : midp / 1e4;
  // auto output_bps = signal / signal_to_bp_divisor;

  // limit payup only
  pie_in_the_sky = std::min<double>(pie_in_the_sky, lm_agg_config_.pie_in_the_sky_cap_bps());
  pie_in_the_sky = std::min<double>(pie_in_the_sky, lm_agg_config_.payup_cap_bps());

  // pie in the sky:
  // model penalty of profit due to predicted missing fills
  if (signal_lb - pie_in_the_sky < leaned_threshold_bps) {
    return;
  }

  /* ---------------------------------- lean ---------------------------------- */
  // if (lean_bps_ > 0) {
  //   double neutralized_potential_pos =
  //   double norm = neutralized_potential_pos / (max_pos - min_pos) * 2
  // }

  /* ------------------------------- agg pricing ------------------------------ */

  double payup_bps =
      signal - threshold_bps_ * lm_agg_config_.threshold_payup_margin_ratio();

  payup_bps = std::min(payup_bps, lm_agg_config_.payup_cap_bps());

  /* --------------------------------- submit --------------------------------- */

  double raw_qty = lholder_.oe->GetConfig(*lholder_.product_holder).qty().max_lot_size();
  bool post_only = false;

  double raw_price = (pie_in_the_sky < 0) ?
      bbo_price + sign * pie_in_the_sky * (midp * 1e-4)
      : bbo_price + sign * payup_bps * (midp * 1e-4);
  double rounded_price =
      lholder_.product_holder->product_info().tick_price().RoundPassively(raw_price, direction);

  if (lm_agg_config_.use_bbo_insteadof_taker()) {
    post_only = true;
    raw_price = bbo_price_pass;
    raw_qty /= 2;
  } else if (pie_in_the_sky < 0) {
    post_only = true;
    raw_qty /= 2;
  } else {
    if (!is_l1_book(*lholder_.last_book)) {
      raw_qty = std::min(
          raw_qty, get_available_qty(*lholder_.last_book, rounded_price, direction));
    }
  }


  if (raw_qty > 0) {
    std::string extra_tag = fmt::format(
        "{{\"logic\": \"lm-agg\", \"model_id\": \"{}\", \"signal\": {}}}",
        model_id,
        signal);
    bool fire_taker = lm_agg_config_.filldepth_window_taker();
    if (direction == OrderDirection::DIRECTION_SELL &&
        prev_sell_taker_time_ + taker_cooltime_ < timestamp_override) {
      prev_sell_taker_time_ = timestamp_override;
    } else if (
        direction == OrderDirection::DIRECTION_BUY &&
        prev_buy_taker_time_ + taker_cooltime_ < timestamp_override) {
      prev_buy_taker_time_ = timestamp_override;
    } else {
      fire_taker = false;
      LOG(INFO) << "Skipped taker due to the cooltime";
    }
    if (fire_taker && !rest_mode_) {
      lholder_.oe->SubmitOrder(
          *lholder_.product_holder,
          rounded_price,  // OE does rounding, so this can be replaced by
                          // ExecutorConfig-RoundingPolicy
          raw_qty,
          direction,
          lm_agg_config_.order_duration(),
          OrderType::LIMIT_ORDER,
          bbo_price,
          leverage_rate_,
          post_only,
          extra_tag,
          timestamp_override,
          false,  // give up logging due to the overhead
          lm_agg_config_.use_block_price_band());
      if (pos_mismatch_calculator_) {
        pos_mismatch_calculator_->UpdateSubmitOrder(timestamp_override, sign, raw_qty);
      }
    }
    LOG(INFO) << extra_tag;
  }
}

}  // namespace coin2::strategy::linear_model
